#include "common.h"
#include "kernel_operator.h"

namespace AscendC
{
template <typename _DT_X, typename _DT_Y> class Kernel
{
  public:
    TPipe pipe;
    CommonTiling tilingData;
    DefInTensor(X);
    DefOutTensor(Y);

    DefBufVECCALC(TEMP1);
    DefBufVECCALC(TEMP2);

  public:
    __aicore__ inline Kernel()
    {
    }
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, CommonTiling &tiling)
    {
        tilingData = tiling;
        // 切分global
        SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP1, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(TEMP2, tilingData.bufferSize * sizeof(_DT_X));
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(i, tilingData.bufferSize);
            CopyOut(i, tilingData.bufferSize);
        }
        auto finalCopyCount = tilingData.getFinalCopyCount<TypeOf(X)>();
        CopyIn(finnal_progress, finalCopyCount);
        Compute(finnal_progress, tilingData.finalCalcCount);
        CopyOut(finnal_progress, finalCopyCount);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(X, index, calcCount);
    }
    // x * e^((x - |x|) * 0.851) / (1 + e^(-1.702 * |x|))
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(X);
        QueAllocSimple(Y);
        TBufGet(TEMP1, _DT_X);
        TBufGet(TEMP2, _DT_X);
        // |x|
        Abs(TBufTensorName(TEMP1), LTensorName(X), calcCount);
        // -1.702 * |x|
        Muls(TBufTensorName(TEMP2), TBufTensorName(TEMP1), _DT_X(-1.702), calcCount);
        // e^(-1.702 * |x|)
        Exp(LTensorName(Y), TBufTensorName(TEMP2), calcCount);
        // e^(-1.702 * |x|) + 1
        Adds(TBufTensorName(TEMP2), LTensorName(Y), _DT_X(1), calcCount);
        // x - |x|
        Sub(LTensorName(Y), LTensorName(X), TBufTensorName(TEMP1), calcCount);
        // (x - |x|) * 0.851
        Muls(TBufTensorName(TEMP1), LTensorName(Y), _DT_X(0.851), calcCount);
        // e^((x - |x|) * 0.851)
        Exp(LTensorName(Y), TBufTensorName(TEMP1), calcCount);
        // e^((x - |x|) * 0.851) * x
        Mul(TBufTensorName(TEMP1), LTensorName(Y), LTensorName(X), calcCount);
        // x * e^((x - |x|) * 0.851) / (1 + e^(-1.702 * |x|))
        Div(LTensorName(Y), TBufTensorName(TEMP1), TBufTensorName(TEMP2), calcCount);
        EnQue(Y);
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount)
    }
};
} // namespace AscendC
extern "C" __global__ __aicore__ void fast_gelu(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    AscendC::Kernel<DTYPE_X, DTYPE_Y> op;
    op.Init(x, y, commonTiling);
    op.Process();
}